package com.wyr.admin.mq.config;

import com.alibaba.fastjson2.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.amqp.AmqpException;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.core.MessagePostProcessor;


/**
 * MQ发消息前添加traceId到消息头上
 *
 * @author wangyiren
 * @version 1.0
 * &#064;title RabbitMessagePostProcessor
 * &#064;description
 * &#064;date 2025/03/06 14:37
 */
@Slf4j
public class RabbitMessagePostProcessor implements MessagePostProcessor {

    /**
     * 日志追踪id
     */
    public static final String TRACE_ID = "traceId";

    @Override
    public Message postProcessMessage(Message message) throws AmqpException {
        // 从消息体中获取
        String traceId = JSONObject.parseObject(new String(message.getBody())).getString(TRACE_ID);
        if (traceId != null) {
            message.getMessageProperties().setHeader(TRACE_ID, traceId);
            MDC.put(TRACE_ID, traceId);
        } else {
            log.warn("MDC does not contain traceId, message may not be traceable");
        }
        return message;
    }

    @Override
    public Message postProcessMessage(Message message, org.springframework.amqp.core.Correlation correlation) {
        return postProcessMessage(message);
    }
}
